import time
from log import log
from typing import Any, Callable

from starlette.concurrency import iterate_in_threadpool
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.requests import Request
from starlette.responses import Response
from starlette.types import ASGIApp


ECHO_HEADERS = []
NO_ECHO_LOG_PATH = []


class RequestContextLogMiddleware2(BaseHTTPMiddleware):
    """
    所有的请求都会使用此中间件
    """

    async def set_body(self, request: Request):
        receive_ = await request._receive()

        async def receive():
            return receive_

        request._receive = receive

    async def dispatch(
            self, request: Request, call_next: RequestResponseEndpoint
    ) -> Response:
        start_time = time.time()
        await self.set_body(request)
        body = await request.body()
        try:
            body = await request.json()
        except Exception as err:
            pass
        log.info(f" start request [request id = {request.headers.get('x-request-id', '')}] path={request.url.path}")
        log.info("request:{0}, body:{1}".format(request, body))
        response = await call_next(request)
        if response.headers.get('content-type') in ECHO_HEADERS \
                and request.url.path not in NO_ECHO_LOG_PATH:
            response_body = [chunk async for chunk in response.body_iterator]
            response.body_iterator = iterate_in_threadpool(iter(response_body))
            log.info(f"The request {request.url.path}:response_body={response_body[0].decode()}")
        process_time = (time.time() - start_time) * 1000
        formatted_process_time = '{0:.2f}'.format(process_time)
        log.info(f"The request [request id = {request.headers.get('x-request-id', '')}] completed_in="
                 f"{formatted_process_time}ms status_code={response.status_code}")
        return response


class RequestContextLogMiddleware:
    def __init__(self, app: ASGIApp) -> None:
        self.app = app

    async def __call__(self, request: Request, call_next: Callable[[Request], Any]) -> Response:
        # 在请求处理前记录请求信息
        self.log_request_context(request)
        response = await call_next(request)

        # 在请求处理后记录响应信息
        self.log_response_context(response)
        return response

    def log_request_context(self, request: Request) -> None:
        # 记录请求方法、URL和客户端IP
        log.info(f"Request Method: {request.method} URL: {request.url} Client IP: {request.client.host}")

    def log_response_context(self, response: Response) -> None:
        # 记录响应状态码和内容长度
        log.info(f"Response Status Code: {response.status_code} Content Length: {len(response.body)}")
